// Copyright Materialize, Inc. All rights reserved.
//
// Use of this software is governed by the Business Source License
// included in the LICENSE file.
//
// As of the Change Date specified in that file, in accordance with
// the Business Source License, use of this software will be governed
// by the Apache License, Version 2.0.

//! Multiple-producer, multiple-consumer channels.
//!
//! Unlike [MPSC channels][mpsc], broadcast channels must be statically
//! allocated. Every broadcast channel must have a unique type that implements
//! [`Token`][broadcast::Token]. Access to a particular broadcast channel is
//! controlled by the visibility of this `Token`, as the token must be presented
//! to the [`Switchboard`] to construct broadcast transmitters or receivers.
//!
//! Broadcast channels always broadcast messages to all nodes in the cluster,
//! with the exception that the broadcasting node can choose whether it would
//! like to receive its own messages, via
//! [`Token::loopback`][broadcast::Token::loopback]. It is therefore very
//! important that all nodes in the cluster allocate their broadcast receiver
//! and periodically drain it of messages; otherwise, the network buffer will
//! fill up, and the transmitter will be unable to broadcast new messages.
//!
//! A relatively contrived example of broadcasting in a one-node cluster
//! follows:
//!
//! ```
//! use comm::{broadcast, Switchboard};
//! use futures::sink::SinkExt;
//! use futures::stream::StreamExt;
//!
//! struct UniversalAnswersToken;
//!
//! impl broadcast::Token for UniversalAnswersToken {
//!     type Item = usize;
//!
//!     fn loopback(&self) -> bool {
//!         // Enable loopback so that we receive our own transmissions.
//!         true
//!     }
//! }
//!
//! let (switchboard, mut runtime) = Switchboard::local()?;
//! let mut tx = switchboard.broadcast_tx(UniversalAnswersToken);
//! let mut rx = switchboard.broadcast_rx(UniversalAnswersToken);
//! runtime.block_on(tx.send(42))?;
//! assert_eq!(runtime.block_on(rx.next()).transpose()?, Some(42));
//! # Ok::<(), Box<dyn std::error::Error>>(())
//! ```

use std::fmt;
use std::future::Future;
use std::pin::Pin;
use std::task::{Context, Poll};

use futures::future::{self, BoxFuture, TryFutureExt};
use futures::ready;
use futures::sink::{Sink, SinkExt};
use futures::stream::{FuturesUnordered, Stream, StreamExt, TryStreamExt};
use serde::{Deserialize, Serialize};
use uuid::Uuid;

use crate::error::Error;
use crate::mpsc;
use crate::protocol::{self, SendSink};
use crate::switchboard::Switchboard;

/// The capability to construct a particular broadcast sender or receiver.
pub trait Token {
    /// The type of the items that will be sent along the channel.
    type Item: Serialize + for<'de> Deserialize<'de> + Send + Unpin + Clone;

    /// Whether transmissions on this channel should be sent only to peers or
    /// sent also to the receiver on the broadcasting node. By default, loopback
    /// is disabled.
    fn loopback(&self) -> bool {
        false
    }

    /// Generates the channel UUID associated with this `Token`. By default,
    /// this derives a UUID from the type's ID, but the method can be overridden
    /// so that the UUID is derived from other properties of the `Token`.
    ///
    /// The implementation must guarantee that the UUID generation process does
    /// not cause collisions with other types that implement `Token`, or with
    /// UUIDs generated by [`Uuid::new_v4`].
    fn uuid(&self) -> Uuid
    where
        Self: 'static,
    {
        use std::any::TypeId;
        use std::collections::hash_map::DefaultHasher;
        use std::hash::{Hash, Hasher};

        let mut hasher = DefaultHasher::new();
        TypeId::of::<Self>().hash(&mut hasher);
        let mut buf = [0; 16];
        (&mut buf[8..]).copy_from_slice(&hasher.finish().to_be_bytes());
        Uuid::from_bytes(buf)
    }
}

/// The transmission end of a broadcast channel.
///
/// See the [`futures::Sink`] documentation for details about the API for
/// sending messages to a sink.
pub struct Sender<D> {
    state: SenderState<D>,
}

enum SenderState<D> {
    Connecting(BoxFuture<'static, Result<SendSink<D>, bincode::Error>>),
    Ready(SendSink<D>),
}

impl<D> fmt::Debug for SenderState<D> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match self {
            SenderState::Connecting { .. } => f.write_str("SenderState::Connecting"),
            SenderState::Ready(_) => f.write_str("SenderState::Ready"),
        }
    }
}

impl<D> Sender<D>
where
    D: Serialize + for<'de> Deserialize<'de> + Send + Unpin + Clone + 'static,
{
    pub(crate) fn new<C, I>(uuid: Uuid, addrs: I) -> Sender<D>
    where
        C: protocol::Connection,
        I: IntoIterator<Item = C::Addr>,
    {
        let conns = addrs
            .into_iter()
            .map(|addr| protocol::connect_channel::<C, D>(addr, uuid))
            .collect::<FuturesUnordered<_>>()
            // TODO(benesch): this might be more efficient with a multi-fanout that
            // could fan out to multiple streams at once. Not clear what the
            // performance of this binary tree of fanouts is.
            .try_fold(
                Box::pin(ore::future::dev_null()) as SendSink<D>,
                |memo, sink| future::ok(Box::pin(memo.fanout(sink)) as SendSink<D>),
            )
            .err_into();

        Sender {
            state: SenderState::Connecting(Box::pin(conns)),
        }
    }
}

impl<D> Sink<D> for Sender<D> {
    type Error = Error;

    fn poll_ready(mut self: Pin<&mut Self>, cx: &mut Context) -> Poll<Result<(), Self::Error>> {
        match &mut self.state {
            SenderState::Connecting(fut) => {
                let sink = ready!(Pin::new(fut).poll(cx))?;
                self.state = SenderState::Ready(sink);
                Poll::Ready(Ok(()))
            }
            SenderState::Ready(sink) => Pin::new(sink).poll_ready(cx),
        }
    }

    fn start_send(mut self: Pin<&mut Self>, item: D) -> Result<(), Self::Error> {
        match &mut self.state {
            SenderState::Connecting(_) => panic!("sending to not ready sink"),
            SenderState::Ready(sink) => Pin::new(sink).start_send(item),
        }
    }

    fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context) -> Poll<Result<(), Self::Error>> {
        match &mut self.state {
            SenderState::Connecting(_) => Poll::Ready(Ok(())),
            SenderState::Ready(sink) => Pin::new(sink).poll_flush(cx),
        }
    }

    fn poll_close(mut self: Pin<&mut Self>, cx: &mut Context) -> Poll<Result<(), Self::Error>> {
        match &mut self.state {
            SenderState::Connecting(_) => self.poll_ready(cx),
            SenderState::Ready(sink) => Pin::new(sink).poll_close(cx),
        }
    }
}

/// The receiving end of a broadcast channel.
///
/// See the [`futures::Stream`] documentation for details about the API for
/// receiving messages from a stream.
pub struct Receiver<D>(mpsc::Receiver<D>);

impl<D> Receiver<D>
where
    D: Serialize + for<'de> Deserialize<'de> + Send + Unpin + 'static,
{
    pub(crate) fn new<C>(
        conn_rx: impl Stream<Item = protocol::Framed<C>> + Send + Unpin + 'static,
        switchboard: Switchboard<C>,
    ) -> Receiver<D>
    where
        C: protocol::Connection,
    {
        Receiver(mpsc::Receiver::new(conn_rx, switchboard, None))
    }

    /// Arranges to split the receiver into multiple receivers, e.g., so that
    /// multiple threads can receive the transmission stream. All receivers will
    /// receive all messages, and will receive them in the same order.
    pub fn fanout(self) -> FanoutReceiverBuilder<D> {
        FanoutReceiverBuilder {
            stream: self,
            sink: Some(Box::pin(ore::future::dev_null())),
        }
    }
}

impl<D> Stream for Receiver<D> {
    type Item = Result<D, Error>;

    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context) -> Poll<Option<Self::Item>> {
        Pin::new(&mut self.0).poll_next(cx)
    }
}

/// The builder returned by [`Receiver::fanout`]. New receivers can be allocated
/// with [`FanoutReceiverBuilder::attach`]. Once all new receivers are
/// allocated, [`FanoutReceiverBuilder::shuttle`] must be called, and the
/// returned future spawned onto a [`tokio::executor::Executor`], in order for
/// the transmissions to be routed appropriately.
pub struct FanoutReceiverBuilder<D> {
    stream: Receiver<D>,
    sink: Option<Pin<Box<dyn Sink<D, Error = FanoutError> + Send>>>,
}

impl<D> FanoutReceiverBuilder<D>
where
    D: Send + Unpin + Clone + 'static,
{
    /// Acquires a new receiver.
    pub fn attach(&mut self) -> futures::channel::mpsc::UnboundedReceiver<D> {
        let (tx, rx) = futures::channel::mpsc::unbounded();
        // TODO(benesch): this might be more efficient with a multi-fanout that
        // could fan out to multiple streams at once. Not clear what the
        // performance of this binary tree of fanouts is.
        self.sink = Some(Box::pin(
            self.sink.take().unwrap().fanout(tx.sink_err_into()),
        ));
        rx
    }

    /// Consumes the builder and returns a future that will shuttle messages
    /// to the fanned-out receivers until the channel is closed.
    pub fn shuttle(self) -> impl Future<Output = Result<(), FanoutError>> {
        self.stream.err_into().forward(self.sink.unwrap())
    }
}

/// The error returned while shuttling messages to fanned-out broadcast
/// receivers.
#[derive(Debug)]
pub enum FanoutError {
    /// An error occurred while receiving messages from upstream.
    ReceiveError(Error),
    /// An error occurred while sending messages downstream.
    SendError(futures::channel::mpsc::SendError),
}

impl fmt::Display for FanoutError {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match self {
            FanoutError::ReceiveError(err) => write!(f, "fanout error: receive error: {}", err),
            FanoutError::SendError(err) => write!(f, "fanout error: send error: {}", err),
        }
    }
}

impl From<Error> for FanoutError {
    fn from(err: Error) -> FanoutError {
        FanoutError::ReceiveError(err)
    }
}

impl From<futures::channel::mpsc::SendError> for FanoutError {
    fn from(err: futures::channel::mpsc::SendError) -> FanoutError {
        FanoutError::SendError(err)
    }
}

impl std::error::Error for FanoutError {}
